#include <TStopwatch.h>
#include <CpLog.h>

#include "CpFitter.h"
#include "CpPolynom1D.h"
#include "CpPolynom2D.h"
#include "CpPolynom3D.h"
#include <CpNames.h>

ClassImp(CpFitter);

//______________________________________________________________________________
CpFitter::CpFitter(): TObject(),
   fOutFile(0),
   fOutputDirName("/tmp"),
   fOutputFileName("defaultOutFile.root"),
   fPolynomGradeNumber(1),
   fMagneticAxisIndex(0)
{
   CpDebug(CpLog::kDebug + 5, "<-");
   Int_t i;

   for (i = 0; i < 3; i++) {
      fMinMaxBin[i][0] = 0.0;
      fMinMaxBin[i][1] = 1.0;
      fMinMaxBin[i][2] = 0.0;
   }

   SetActivePositionAxis();
   CpDebug(CpLog::kDebug + 5, "->");
}

//______________________________________________________________________________
CpFitter::~CpFitter()
{
   CpDebug(CpLog::kDebug + 5, "<-");
   CpDebug(CpLog::kDebug + 5, "->");
}



void CpFitter::Print(Option_t *)
{
   CpInfo(Form("Polynom Number                   : %d", fPolynomGradeNumber));
   CpInfo(Form("Number Of Active Postition Axis  : %d", GetNumberOfPositionAxis()));
   CpInfo(Form("Magentic Axis Index              : %d", fMagneticAxisIndex));
   CpInfo(Form("Axis X                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[0], fMinMaxBin[0][0], fMinMaxBin[0][1], fMinMaxBin[0][2]));
   CpInfo(Form("Axis Y                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[1], fMinMaxBin[1][0], fMinMaxBin[1][1], fMinMaxBin[1][2]));
   CpInfo(Form("Axis Z                           : %d [%.3f,%.3f] with %.0f bins", fActivePositionAxis[2], fMinMaxBin[2][0], fMinMaxBin[2][1], fMinMaxBin[2][2]));
   CpInfo("");
   CpInfo(Form("Output file is %s/%s", fOutputDirName.Data(), fOutputFileName.Data()));
   CpInfo("");

}


Bool_t CpFitter::Run(CpMagneticFieldMap *fieldMap)
{

   // Prints current setting
   Print();

   // TODO outputfile (from config)
   fOutFile = TFile::Open(Form("%s/%s", fOutputDirName.Data(), fOutputFileName.Data()), "RECREATE");

   switch (GetNumberOfPositionAxis()) {
      case 1:
         Fit1D(fieldMap);
         break;
      case 2:
         Fit2D(fieldMap);
         break;
      case 3:
         Fit3D(fieldMap);
         break;
      default:
         CpError("Number of position axis is different then 1. Skipping ...");
         break;
   }

   // close file and set pointer to 0
   delete fOutFile;
   fOutFile = 0;

   // return true
   return kTRUE;
}


void CpFitter::Fit1D(CpMagneticFieldMap *fieldMap)
{

   CpInfo("Fitting 1D ...");
   TStopwatch timer;
   timer.Start();
   Double_t min[3];
   Double_t max[3];
   Int_t nBins[3];
   Double_t constNum[3];
   Int_t constNumIndex[3];
   Int_t activeAxisIndex[] = {0, 0, 0};


   Int_t numMinMaxBinNum = 0;
   Int_t numConstNum = 0;

   for (Int_t i = 0; i < 3; i++) {
      if (fActivePositionAxis[i]) {
         min[numMinMaxBinNum] = fMinMaxBin[i][0];
         max[numMinMaxBinNum] = fMinMaxBin[i][1];
         nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
         activeAxisIndex[numMinMaxBinNum] = i;
         numMinMaxBinNum++;
      } else {
         constNum[numConstNum] = fMinMaxBin[i][0];
         constNumIndex[numConstNum] = i;
         numConstNum++;
      }
   }

   TString polynomName(Form("fun_%s_1D", CpNames::AxisName(activeAxisIndex[0])));
   TF1 *myPol_1D = GetFittingFunction1D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0]);

   if (!myPol_1D) {
      CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
      return;
   }

   CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

   TH1D *hist1D = new TH1D(Form("hist_%s_1D", CpNames::AxisName(activeAxisIndex[0])), Form("H_{%s} for %s=%.2f, %s=%.2f", CpNames::AxisName(fMagneticAxisIndex), CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]), nBins[0], min[0], max[0]);
   fieldMap->FillHistogram(hist1D, fMagneticAxisIndex, activeAxisIndex[0], constNum[0], constNum[1]);
   hist1D->Fit(myPol_1D, "QN");


   if (!fOutFile) {
      CpWarning("fOutFile is null !!!");
   } else {
      CpInfo("=================================================");
      CpInfo("\t\tFitting results");
      CpInfo("=================================================");
      CpInfo("Data histogram");
      hist1D->Print();
      CpInfo("");
      CpInfo("Fit function");
      myPol_1D->Print();
      CpInfo("=================================================");
      fOutFile->cd();
      hist1D->Write();
      myPol_1D->Write();
   }

   timer.Stop();
   timer.Print();
}

void CpFitter::Fit2D(CpMagneticFieldMap *fieldMap)
{

   CpInfo("Fitting 2D ...");
   TStopwatch timer;
   timer.Start();
   Double_t min[3];
   Double_t max[3];
   Int_t nBins[3];
   Double_t constNum[3];
   Int_t constNumIndex[3];
   Int_t activeAxisIndex[] = {0, 0, 0};


   Int_t numMinMaxBinNum = 0;
   Int_t numConstNum = 0;

   for (Int_t i = 0; i < 3; i++) {
      if (fActivePositionAxis[i]) {
         min[numMinMaxBinNum] = fMinMaxBin[i][0];
         max[numMinMaxBinNum] = fMinMaxBin[i][1];
         nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
         activeAxisIndex[numMinMaxBinNum] = i;
         numMinMaxBinNum++;
      } else {
         constNum[numConstNum] = fMinMaxBin[i][0];
         constNumIndex[numConstNum] = i;
         numConstNum++;
      }
   }

   TString polynomName(Form("fun_%s%s_2D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1])));
   TF2 *myPol_2D = GetFittingFunction2D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0], min[1], max[1]);

   if (!myPol_2D) {
      CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
      return;
   }

//   CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

   TH2D *hist2D = new TH2D(Form("hist_%s%s_2D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1])), Form("H_{%s} for %s=%.2f", CpNames::AxisName(fMagneticAxisIndex), CpNames::AxisName(constNumIndex[0]), constNum[0]), nBins[0], min[0], max[0], nBins[1], min[1], max[1]);
   fieldMap->FillHistogram(hist2D, fMagneticAxisIndex, activeAxisIndex[0], activeAxisIndex[1], constNum[0]);
   hist2D->Fit(myPol_2D, "QN");

   if (!fOutFile) {
      CpWarning("fOutFile is null !!!");
   } else {
      CpInfo("=================================================");
      CpInfo("\t\tFitting results");
      CpInfo("=================================================");
      CpInfo("Data histogram");
      hist2D->Print();
      CpInfo("");
      CpInfo("Fit function");
      myPol_2D->Print();
      CpInfo("=================================================");
      fOutFile->cd();
      hist2D->Write();
      myPol_2D->Write();
   }

   timer.Stop();
   timer.Print();
}

void CpFitter::Fit3D(CpMagneticFieldMap *fieldMap)
{

   CpInfo("Fitting 3D ...");
   TStopwatch timer;
   timer.Start();
   Double_t min[3];
   Double_t max[3];
   Int_t nBins[3];
   Double_t constNum[3];
   Int_t constNumIndex[3];
   Int_t activeAxisIndex[] = {0, 0, 0};

   Int_t numMinMaxBinNum = 0;
   Int_t numConstNum = 0;

   for (Int_t i = 0; i < 3; i++) {
      if (fActivePositionAxis[i]) {
         min[numMinMaxBinNum] = fMinMaxBin[i][0];
         max[numMinMaxBinNum] = fMinMaxBin[i][1];
         nBins[numMinMaxBinNum] = (Int_t) fMinMaxBin[i][2];
         activeAxisIndex[numMinMaxBinNum] = i;
         numMinMaxBinNum++;
      } else {
         constNum[numConstNum] = fMinMaxBin[i][0];
         constNumIndex[numConstNum] = i;
         numConstNum++;
      }
   }

   TString polynomName(Form("fun_%s%s%s_3D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1]), CpNames::AxisName(activeAxisIndex[2])));
   TF3 *myPol_3D = GetFittingFunction3D(polynomName.Data(), fPolynomGradeNumber, min[0], max[0], min[1], max[1], min[2], max[2]);

   if (!myPol_3D) {
      CpError(Form("Function %s is null. Returning ... ", polynomName.Data()));
      return;
   }

//   CpDebug(CpLog::kDebug, Form("constNum[%s]=%f, constNum[%s]=%f", CpNames::AxisName(constNumIndex[0]), constNum[0], CpNames::AxisName(constNumIndex[1]), constNum[1]));

   TH3D *hist3D = new TH3D(Form("hist_%s%s%s_3D", CpNames::AxisName(activeAxisIndex[0]), CpNames::AxisName(activeAxisIndex[1]), CpNames::AxisName(activeAxisIndex[2])), Form("H_{%s}", CpNames::AxisName(fMagneticAxisIndex)), nBins[0], min[0], max[0], nBins[1], min[1], max[1], nBins[2], min[2], max[2]);
   fieldMap->FillHistogram(hist3D, fMagneticAxisIndex);
   hist3D->Fit(myPol_3D, "QN");

   if (!fOutFile) {
      CpWarning("fOutFile is null !!!");
   } else {
      CpInfo("=================================================");
      CpInfo("\t\tFitting results");
      CpInfo("=================================================");
      CpInfo("Data histogram");
      hist3D->Print();
      CpInfo("");
      CpInfo("Fit function");
      myPol_3D->Print();
      CpInfo("=================================================");
      fOutFile->cd();
      hist3D->Write();
      myPol_3D->Write();
   }

   timer.Stop();
   timer.Print();
}

TF1 *CpFitter::GetFittingFunction1D(TString name, Int_t polynomNum, Double_t min, Double_t max)
{
   TF1 *myPol = 0;
   Int_t nParam = CpPolynom1D::GetNumberOfParameters(polynomNum);
   CpInfo(Form("Number of parameters : %d", nParam));

   switch (polynomNum) {
      case 0:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree0, min, max, nParam);
         break;
      case 1:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree1, min, max, nParam);
         break;
      case 2:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree2, min, max, nParam);
         break;
      case 3:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree3, min, max, nParam);
         break;
      case 4:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree4, min, max, nParam);
         break;
      case 5:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree5, min, max, nParam);
         break;
      case 6:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree6, min, max, nParam);
         break;
      case 7:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree7, min, max, nParam);
         break;
      case 8:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree8, min, max, nParam);
         break;
      case 9:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree9, min, max, nParam);
         break;
      case 10:
         myPol = new TF1(name.Data(), &CpPolynom1D::Degree10, min, max, nParam);
         break;
      default:
         break;
   }

   return myPol;
}

TF2 *CpFitter::GetFittingFunction2D(TString name, Int_t polynomNum, Double_t min1, Double_t max1, Double_t min2, Double_t max2)
{
   TF2 *myPol = 0;
   Int_t nParam = CpPolynom2D::GetNumberOfParameters(polynomNum);
   CpInfo(Form("Number of parameters : %d", nParam));

   switch (polynomNum) {
      case 0:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree0, min1, max1, min2, max2, nParam);
         break;
      case 1:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree1, min1, max1, min2, max2, nParam);
         break;
      case 2:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree2, min1, max1, min2, max2, nParam);
         break;
      case 3:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree3, min1, max1, min2, max2, nParam);
         break;
      case 4:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree4, min1, max1, min2, max2, nParam);
         break;
      case 5:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree5, min1, max1, min2, max2, nParam);
         break;
      case 6:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree6, min1, max1, min2, max2, nParam);
         break;
      case 7:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree7, min1, max1, min2, max2, nParam);
         break;
      case 8:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree8, min1, max1, min2, max2, nParam);
         break;
      case 9:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree9, min1, max1, min2, max2, nParam);
         break;
      case 10:
         myPol = new TF2(name.Data(), &CpPolynom2D::Degree10, min1, max1, min2, max2, nParam);
         break;
      default:
         break;
   }

   return myPol;
}

TF3 *CpFitter::GetFittingFunction3D(TString name, Int_t polynomNum, Double_t min1, Double_t max1, Double_t min2, Double_t max2, Double_t min3, Double_t max3)
{
   TF3 *myPol = 0;
   Int_t nParam = CpPolynom3D::GetNumberOfParameters(polynomNum);
   CpInfo(Form("Number of parameters : %d", nParam));

   switch (polynomNum) {
      case 0:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree0, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 1:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree1, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 2:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree2, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 3:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree3, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 4:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree4, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 5:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree5, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 6:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree6, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 7:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree7, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 8:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree8, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 9:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree9, min1, max1, min2, max2, min3, max3, nParam);
         break;
      case 10:
         myPol = new TF3(name.Data(), &CpPolynom3D::Degree10, min1, max1, min2, max2, min3, max3, nParam);
         break;
      default:
         break;
   }

   return myPol;
}

void CpFitter::SetMinMaxBin(Int_t index, Double_t min, Double_t max, Int_t bins)
{
   CpDebug(CpLog::kDebug + 5, "<-");
   fMinMaxBin[index][0] = min;
   fMinMaxBin[index][1] = max;
   fMinMaxBin[index][2] = bins;
   CpDebug(CpLog::kDebug + 5, "->");
}


void CpFitter::SetActivePositionAxis(Bool_t isX, Bool_t isY, Bool_t isZ)
{
   CpDebug(CpLog::kDebug + 5, "<-");
   fActivePositionAxis[0] = isX;
   fActivePositionAxis[1] = isY;
   fActivePositionAxis[2] = isZ;
   CpDebug(CpLog::kDebug + 5, "->");
}


void CpFitter::SetActivePositionAxis(Int_t index, Bool_t isActive)
{
   fActivePositionAxis[index] = isActive;
}

Int_t CpFitter::GetNumberOfPositionAxis()
{
   CpDebug(CpLog::kDebug + 5, "<-");
   Int_t i, num = 0;

   for (i = 0; i < 3; i++)
      if (fActivePositionAxis[i]) num++;

   CpDebug(CpLog::kDebug + 5, "->");
   return num;

}

void CpFitter::SetPositionsAxisFromString(TString str)
{
   CpDebug(CpLog::kDebug + 5, "<-");
   CpDebug(CpLog::kDebug, Form("%s", str.Data()));
   TObjArray *array = str.Tokenize("][");
   TObjArray *array2 = 0;
   TString strr;

   for (Int_t i = 0; i < array->GetEntriesFast(); i++) {
      strr = ((TObjString *) array->At(i))->GetString();

      if (strr.Contains(";")) {
         SetActivePositionAxis(i, kTRUE);
         array2 = strr.Tokenize(";");
         SetMinMaxBin(i, (((TObjString *) array2->At(0))->GetString()).Atof(), (((TObjString *) array2->At(1))->GetString()).Atof(), (((TObjString *) array2->At(2))->GetString()).Atof());
      } else {
         SetActivePositionAxis(i, kFALSE);
         SetMinMaxBin(i, strr.Atof(), strr.Atof(), 0);

      }

      CpDebug(CpLog::kDebug, Form("%s", strr.Data()));
   }

   CpDebug(CpLog::kDebug + 5, "->");
}
